# from typing_extensions import ParamSpec
from typing import KeysView
from graph_cut import get_hist, get_histx, get_mask, graph_cut, mask_add_depth

import argparse
# import torch
# import torch.nn as nn
# from torch.utils.data import DataLoader
from tqdm import tqdm
import numpy as np
import math
import os
import json
import cv2
# torch.backends.cudnn.enabled = False
import matplotlib.pyplot as plt


parser = argparse.ArgumentParser(description='Greenhouse Image Analysis')
parser.add_argument('--image_dir',   type=str,   default='C:\\Users\\ys\\Desktop\\data_original\\', help='image_dir')
parser.add_argument('--gt_file',     type=str,   default='GroundTruth_All_388_Images.json', help='gt_file')
# parser.add_argument('--output_dir',   type=str,   default='./output_whitebg/', help='output_dir')
parser.add_argument('--variety',   type=str,   default='Salanova', help='Aphylion, Salanova, Satine, Lugano')
params = parser.parse_args()
print(params)

json_path = os.path.join(params.image_dir, f'GroundTruth_All_388_Images.json')
jsn_f = open(json_path, 'r')
jsn = json.load(jsn_f)
jsn_dataset = jsn['Measurements']

a_size = 15



keys = []
for key, value in jsn_dataset.items():
    if value["Variety"] == params.variety:
        keys.append(key)

for key in tqdm(keys, desc=params.variety):
    value=jsn_dataset[key]
    output_dir = f'./data_graphcut/'
    # output_dir = f'./output/{value["Variety"]}/'
    rgb_filename = value['RGB_Image']
    depth_filename = value['Depth_Information']
    histx = get_histx(value["Variety"])

    # print(value["Variety"])
    img_rgb_path = os.path.join(params.image_dir+'RGBImages', rgb_filename)
    img_depth_path = os.path.join(params.image_dir+'DepthImages', depth_filename)
    img_rgb = cv2.imread(img_rgb_path)
    img_depth = cv2.imread(img_depth_path, cv2.IMREAD_UNCHANGED)
    img_depth = np.where(img_depth == 0, 65535, img_depth)

    
    mask = get_mask(img_rgb, a_size, histx)
    mask = mask_add_depth(img_depth, mask)
    # cv2.imwrite(output_dir+rgb_filename[:-4]+'.jpg', mask*70)

    rgb_result, depth_result, mask = graph_cut(img_rgb, img_depth, mask=mask)

    os.makedirs(os.path.join(output_dir+'RGBImages'), exist_ok=True)
    os.makedirs(os.path.join(output_dir+'DepthImages'), exist_ok=True)
    cv2.imwrite(os.path.join(output_dir, 'RGBImages', rgb_filename), rgb_result)
    cv2.imwrite(os.path.join(output_dir, 'DepthImages', depth_filename), depth_result)
    
    # r = (486, 914, 486+155+50, 914+167+100)
    # img_ori = cv2.rectangle(img_ori, (r[1], r[0]), (r[3], r[2]), (0, 0, 0), 3)
    # cv2.imshow('img', depth_result*200)
    # cv2.waitKey(300)

# cv2.imwrite(f'./output/{filename}', result)


